package com.xs.langchain4j_springboot.controller;

import com.xs.langchain4j_springboot.service.DownLoadFileService;
import com.xs.langchain4j_springboot.config.AiConfig;
import com.xs.langchain4j_springboot.config.PdfAuditAgent;
import com.xs.langchain4j_springboot.service.FileReaderService;
import dev.langchain4j.community.model.dashscope.QwenChatModel;
import dev.langchain4j.community.model.dashscope.QwenStreamingChatModel;
import dev.langchain4j.model.chat.response.ChatResponse;
import dev.langchain4j.model.chat.response.StreamingChatResponseHandler;
import dev.langchain4j.model.image.ImageModel;
import dev.langchain4j.model.output.Response;
import dev.langchain4j.service.TokenStream;
import jakarta.annotation.Resource;
import org.apache.pdfbox.pdmodel.PDDocument;
import org.apache.pdfbox.text.PDFTextStripper;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.web.bind.annotation.RequestMapping;
import org.springframework.web.bind.annotation.RequestParam;
import org.springframework.web.bind.annotation.RestController;
import org.springframework.web.servlet.View;
import reactor.core.publisher.Flux;

import java.io.IOException;
import java.io.InputStream;
import java.net.URL;
import java.time.LocalDate;

@RestController
@RequestMapping("/ai")
public class ChatController {
    @Autowired
    QwenChatModel qwenChatModel;

    @Autowired
    QwenStreamingChatModel qwenStreamingChatModel;
    @Autowired
    private View error;

    @Autowired
    AiConfig.Assistant assistant;

    @Autowired
    PdfAuditAgent.auditAgent auditAgent;

    @Autowired
    AiConfig.AssistantUnique assistantUnique;

    @Resource
    DownLoadFileService downLoadFileService;

    @Resource
    FileReaderService fileReaderService;


    @RequestMapping("/chat")
    public String test(@RequestParam(defaultValue = "你是谁") String message) {
        String chat = qwenChatModel.chat(message);
        return chat;
    }

    @RequestMapping(value = "/stream",produces = "text/stream;charset=UTF-8")
    public Flux<String> streamingTest(@RequestParam(defaultValue = "你是谁") String message) {
        Flux<String> res = Flux.create(fluxSink -> {
            qwenStreamingChatModel.chat(message, new StreamingChatResponseHandler() {
                @Override
                public void onPartialResponse(String s) {
                    fluxSink.next(s);
//                    System.out.println(s);
                }

                @Override
                public void onCompleteResponse(ChatResponse chatResponse) {
                    fluxSink.complete();
                }

                @Override
                public void onError(Throwable throwable) {
                    fluxSink.error(throwable);
                }
            });
        });

        return res;
    }

    @RequestMapping(value = "/memory_chat")
    public String memoryChat(@RequestParam(defaultValue = "我叫胡治纬") String message) {
        return assistant.chat(message);
    }

    @RequestMapping(value = "/error")
    public String error() {
        return "错误";
    }

    @RequestMapping(value = "/category_classify",produces = "text/stream;charset=UTF-8")
    public Flux<String> categoryClassify(@RequestParam String message) {

        TokenStream stream = assistant.categoryAudit(message);
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(s -> fluxSink.next(s))
                    .onCompleteResponse(c -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
    }

    @RequestMapping(value = "/memory_stream_chat",produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryStreamChat(@RequestParam(defaultValue = "我是谁") String message) {
        TokenStream stream = assistant.ticketAssistant(message, LocalDate.now().toString());
//        String something = assistant.doSomething(message, "en");
//        System.out.println(something);
//        String roastMe = assistant.roastMe(message);
//        System.out.println(roastMe);

        return Flux.create(fluxSink -> {
            stream.onPartialResponse(s -> fluxSink.next(s))
                    .onCompleteResponse(c -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });

    }

    @RequestMapping(value = "/memory_unique_chat",produces = "text/stream;charset=UTF-8")
    public Flux<String> memoryUniqueChat(@RequestParam(defaultValue = "我是谁") String message, Integer userId) {
        TokenStream stream = assistantUnique.stream(userId, message);
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(s -> fluxSink.next(s))
                    .onCompleteResponse(c -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
//        return assistantUnique.chat(1, message);
    }

    @RequestMapping(value = "/pdf_audit",produces = "text/stream;charset=UTF-8")
    public Flux<String> pdfAudit(String url) throws IOException, InterruptedException {
        /**
         * 大模型对pdf的阅读理解基本没啥问题，都挺准的，不用下载到本地也可以读取内容并理解
         * 但是对于图片的识别，OCR好像不是很好用，图片可能得下载到本地然后再扔到大模型里取获取内容并解析
         */
        String content = "";
        if (url.endsWith("pdf")) {
            content = fileReaderService.readPdfFile(url);
            System.out.println(content);
        } else if (url.endsWith("jpg") || url.endsWith("png")) {
            /**
             * 直接使用OCR的效果不是很好，解读不完整，准确率也不高，目前来说图片格式的审核文件还是少数，我觉得可以先不做这一块
             * 要么就选择直接把图片下载到本地然后扔给千问大模型去解读
             */
        }
        TokenStream stream = auditAgent.pdfAudit(content);
        return Flux.create(fluxSink -> {
            stream.onPartialResponse(s -> fluxSink.next(s))
                    .onCompleteResponse(c -> fluxSink.complete())
                    .onError(fluxSink::error)
                    .start();
        });
    }

}
